# 四种记忆机制
# 1. ConversationBufferMemory: 缓冲记忆
# 优势: 为LLM提供最大量的信息,方法简单,直观
# 劣势: 使用更多的 Token,导致响应时间增加和更高的成本; 长时间对话可能超出 Token限制

# 2. ConversationBufferWindowMemory: 窗口记忆
# 优势: 只保留最近的互动,Token使用量较少; 窗口大小颗调节,灵活度高
# 劣势: 无法记住早期互动; 太大的窗口就跟 ConversationBufferMemory 类似了

# 3. ConversationSummaryMemory: 总结记忆
# 优势: 允许更长时间对话的同时减少了 Token的使用量
# 劣势: 会存在记忆失真,总结的能力取决于大模型的汇总能力

# 4. ConversationSummaryBufferMemory: 混合记忆
# 优势: 既可以记住早期互动,也可以记住最近互动,同时使用了总结记忆和窗口记忆的优势
# 劣势: 早期内容可能失真,需要根据大模型的总结能力来调节

import os
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory, RedisChatMessageHistory

os.environ["OPENAI_API_KEY"] = "sk-zk2d1745aac506592237476bef6f546af74679cb75331b8a"
BASE_URL = "https://api.zhizengzeng.com/v1/"

# 初始化 OpenAI 的 LLM
llm = ChatOpenAI(
    model_name="gpt-3.5-turbo",
    temperature=0.7,
    base_url=BASE_URL,
)

# 定义对话模板（显式包含历史消息占位符）
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个友好的助手"),  # 可选的系统消息
    MessagesPlaceholder(variable_name="history"),  # 历史消息占位符
    ("human", "{input}"),  # 用户输入占位符
])

# 构建基础对话链（正确传递历史消息）
chain = (
        RunnablePassthrough.assign(
            # 将历史消息转换为消息列表
            history=lambda x: x.get("history", [])
        )
        | prompt
        | llm
)

# 使用内存存储对话历史(可以替换为redis,数据库等),这里是获取对话历史,根据 session_id 来区分不同的客户
get_session_history = lambda session_id: ChatMessageHistory()
# 使用 Redis 存储对话历史,这里是获取对话历史,根据 session_id 来区分不同的客户
# get_session_history = lambda session_id: RedisChatMessageHistory(
#         session_id=session_id,
#         url="redis://localhost:6379"
#     )
#


# 包装为支持历史记录的链
chain_with_history = RunnableWithMessageHistory(
    chain,
    get_session_history,  # 使用内存存储对话历史(可以替换为数据库)
    input_messages_key="input",  # 用户输入的键名
    history_messages_key="history",  # 必须与模板中的 variable_name 一致
)

# 测试对话（使用同一 session_id 保持历史）
response = chain_with_history.invoke(
    {"input": "你好！"},
    config={"configurable": {"session_id": "test_session"}}  # 必须指定 session_id
)
print(response.content)
